try:
    from detect_evaluate import *
except Exception as e:
    print(e)
    from .detect_evaluate import *
import shutil


# 数据格式准备{（json名称）：【类别】}
def prepare_frame(src):
    imgs_classes_dict = {}
    # 每个文件夹对应的json名
    # img_dict = {}
    # for dir in os.listdir(src):
    #     dir_path = os.path.join(src, dir)
    #     img_dict.setdefault(dir, [])
    #     for file in os.listdir(dir_path):
    #         if os.path.splitext(file)[1] == '.json':
    #             img_dict[dir].append(file)
    # print(img_dict)
    # 尝试一起处理，嵌套字典{文件夹名称{json, labels}}
    for dir in os.listdir(src):
        dir_path = os.path.join(src, dir)
        imgs_classes_dict[dir] = {'json_name': [], 'labels': []}
        for file in os.listdir(dir_path):
            if os.path.splitext(file)[1] == '.json':
                imgs_classes_dict[dir]['json_name'].append(file)
            elif os.path.splitext(file)[1] == '.txt':
                with open(os.path.join(dir_path, file)) as f:
                    for label in f.readlines():
                        label = label.replace('\n', '').strip()
                        imgs_classes_dict[dir]['labels'].append(label)
    # 好像有点冗余
    result_dict = {}
    for keys, vals in imgs_classes_dict.items():
        js_tuple = tuple(imgs_classes_dict[keys]['json_name'])
        label_list = imgs_classes_dict[keys]['labels']
        result_dict[js_tuple] = label_list
    # print(result_dict)
    return result_dict


# 2个字典相同的key，列表相加
def dict_same_value_add(dict_a, dict_b):
    # 往a中加入b
    for key, vals in dict_b.items():
        if key in dict_a.keys():
            for val in vals:
                dict_a[key].append(val)
        else:
            dict_a[key] = vals


def main_diff_batch_eval_proc(src_img, res_pt, src_gt, dst_res=None, dst_xls=None):
    # src_img: 打标图片路径
    # res_pt: 打标预测结果json路径
    # src_gt: 人工标注图片、json文件路径
    # dst_res: 将打标结果按src_gt目录进行分类
    # dst_xls: 评估结果保存excel
    result_map = {}
    result_per_img = {}

    json_classes_dict = prepare_frame(src_gt)
    for key, val in json_classes_dict.items():
        one_result_map, one_result_per_img = evaluate_api(src_img, res_pt, list(key), val)
        dict_same_value_add(result_map, one_result_map)
        dict_same_value_add(result_per_img, one_result_per_img)

    if dst_res is not None and dst_xls is None:
        dst_xls = os.path.join(dst_res, 'evaluation.xlsx')

    if dst_res is not None:
        img_ann_list = cvio.load_img_ann_list(src_gt, recursive=True)
        for img, ann in img_ann_list:
            subname = img.split(os.sep)[-2]
            save = os.path.join(dst_res, subname)
            if not os.path.exists(save):
                os.makedirs(save)
            shutil.copy(img, save)
            shutil.copy(ann, save)

    if dst_xls is not None:
        save_to_excel(result_map, result_per_img, dst_xls)


if __name__ == '__main__':
    # result_map = {}
    # result_per_img = {}
    # # excel保存路径
    # save_dst = r'G:\data\datasets\drink\cola\test\20230302\e_top1_0.3_drink_10104_0302-jzx.xlsx'
    # # 装有不同测试集文件夹的根目录
    # json_classes_dict = prepare_frame(
    #     r'G:\data\datasets\drink\cola\test\20230302\gt_dir')
    # for key, val in json_classes_dict.items():
    #     # gtsrc, detsrc, js_list, save_dir, classes, iou_thr=0.5, mode='bbox', iou_mode='standard'
    #     # 参数1：所有测试集大图放到同一个文件夹下的目录   参数2：所有打标的保存路径
    #     one_result_map, one_result_per_img = evaluate_api(
    #         r'G:\data\datasets\drink\cola\test\20230302\gt_all_0302_new',
    #         r'G:\data\datasets\drink\cola\test\20230302\gt_all_0302_new_pred_10104_map',
    #         list(key), val)
    #     dict_same_value_add(result_map, one_result_map)
    #     dict_same_value_add(result_per_img, one_result_per_img)

    # save_to_excel(result_map, result_per_img, save_dst)

    src_img = r'G:\data\datasets\drink\cola\test\20230302\gt_all_0302_new'
    res_pt = r'G:\data\datasets\drink\cola\test\20230302\gt_all_0302_new_pred_10104_map'
    src_gt = r'G:\data\datasets\drink\cola\test\20230302\gt_dir'
    dst_res = r'G:\data\datasets\drink\cola\test\20230302\e_top1_0.3_drink_10104_0302-jzx'
    dst_xls = r'G:\data\datasets\drink\cola\test\20230302\e_top1_0.3_drink_10104_0302-jzx.xlsx'
    main_diff_batch_eval_proc(src_img, res_pt, src_gt, dst_res=dst_res, dst_xls=dst_xls)
